# load data
source("HNC_metadata_tidy.R")

# mutation burdebs
SBSburden <- read.csv("Input_mutational_burden/Input_SBS_mutational_burden.csv") %>% rename(SBS_burden = Mutational.burden)
DBSburden <- read.csv("Input_mutational_burden/Input_DBS_mutational_burden.csv") %>% rename(DBS_burden = Mutational.burden)
IDburden <- read.csv("Input_mutational_burden/Input_ID_mutational_burden.csv") %>% rename(ID_burden = Mutational.burden)

# signature attributions
sbsCOSMIC <- read.csv("Input_signature_attributions/Input_SBS_COSMIC_attributions.csv", stringsAsFactors = T, row.names = 1) %>% tibble::rownames_to_column("donor_id")
dbsCOSMIC <- read.csv("Input_signature_attributions/Input_DBS_COSMIC_attributions.csv", stringsAsFactors = T, row.names = 1) %>% tibble::rownames_to_column("donor_id")
idCOSMIC <- read.csv("Input_signature_attributions/Input_ID_COSMIC_attributions.csv", stringsAsFactors = T, row.names = 1) %>% tibble::rownames_to_column("donor_id")

additional_df <- list(SBSburden, DBSburden, IDburden, sbsCOSMIC, dbsCOSMIC, idCOSMIC)

cosmic <- additional_df[[1]]
for (n in 2:length(additional_df)) {
  cosmic <- cosmic %>% left_join(additional_df[[n]], by = "donor_id")
}

# merge with cosmic sigantures and tidy epi variables
data <- data %>%
  left_join(cosmic, by = "donor_id") %>%
  mutate(
    cigpackyears = ifelse(tobacco_ever == "No", 0, cigpackyears),
    cigqty = ifelse(tobacco_ever == "No", 0, cigqty),
    cigdur = ifelse(tobacco_ever == "No", 0, cigdur),
    age = as.numeric(scale(age_diag))
  )

outcome_var <- c("SBS_burden", "DBS_burden", "ID_burden", "SBS4", "SBS92", "SBS_I", "ID3", "DBS2", "DBS6")
var <- c("cigqty", "cigdur")
confounders <- c("age", "subsite", "sex", "region", "alcohol_ever")

SBS_limit <- 100000
DBS_limit <- 6000
ID_limit <- 5000

result <- NULL
for (v in var) {
  for (sigtype in outcome_var) {
    # remove hypermutated for models including total burdens and variables with missing data
    if (sigtype == "SBS_burden") {
      dat <- data %>% filter(SBS_burden < SBS_limit, !is.na(v))
    } else if (sigtype == "DBS_burden") {
      dat <- data %>% filter(DBS_burden < DBS_limit, !is.na(v))
    } else if (sigtype == "ID_burden") {
      dat <- data %>% filter(ID_burden < ID_limit, !is.na(v))
    } else {
      dat <- data %>% filter(!is.na(cigqty))
    }
    # run linear regression
    myformula <- paste0(sigtype, " ~ ", paste(c(v, confounders), collapse = " + "))
    print(myformula)
    model <- lm(myformula, data = dat)
    r <- summary(model)$coefficients
    r <- r %>%
      as.data.frame() %>%
      mutate(signature = sigtype, .before = everything()) %>%
      tibble::rownames_to_column("independent_vars") %>%
      filter(independent_vars == v)
    result <- rbind(result, r)
  }
}

write.csv(result, "output/Supplementary_Note_Table2.csv", row.names = F)
